import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def _make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v

class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        return F.relu6(x + 3., inplace=self.inplace) / 6.


class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        return x * F.relu6(x + 3., inplace=self.inplace) / 6.


class SELayer(nn.Module):
    def __init__(self, channel, reduction=4):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
                nn.Linear(channel, _make_divisible(channel // reduction, 8)),
                nn.ReLU(inplace=True),
                nn.Linear(_make_divisible(channel // reduction, 8), channel),
                h_sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y

class eca_block(nn.Module):
    def __init__(self, channel, b=1, gamma=2):
        super(eca_block, self).__init__()
        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y)
        return x * y.expand_as(x)

class spacial_attention(nn.Module):
    def __init__(self, kernel_size=7):
        super(spacial_attention, self).__init__()
        padding = kernel_size // 2
        self.conv = nn.Conv2d(1, 1, kernel_size, 1, padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_pool_out    = torch.mean(x, dim=1, keepdim=True)
        out = self.conv(avg_pool_out)
        out = self.sigmoid(out)
        return x * out

def conv_3x3_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        h_swish()
    )


def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        h_swish()
    )


class InvertedResidual(nn.Module):
    def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs, i):
        super(InvertedResidual, self).__init__()
        assert stride in [1, 2]

        self.identity = stride == 1 and inp == oup

        if i <= 1:
            attention = spacial_attention()
        else:
            attention = eca_block(hidden_dim)

        if inp == hidden_dim:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                h_swish() if use_hs == 'HS' else nn.ReLU(inplace=True),
                # Squeeze-and-Excite
                attention if use_se else nn.Identity(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(

                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                h_swish() if use_hs == 'HS' else nn.ReLU(inplace=True),

                # dw
                nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),

                # Squeeze-and-Excite
                attention if use_se else nn.Identity(),

                h_swish() if use_hs == 'HS' else nn.ReLU(inplace=True),

                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.identity:
            return x + self.conv(x)
        else:
            return self.conv(x)

class FusedInvertedResidual(nn.Module):
    def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
        super(FusedInvertedResidual, self).__init__()
        assert stride in [1, 2]

        self.identity = stride == 1 and inp == oup

        self.conv = nn.Sequential(
            # dw
            nn.Conv2d(inp, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2, bias=False),
            nn.BatchNorm2d(hidden_dim),
            h_swish() if use_hs == 'HS' else nn.ReLU(inplace=True),
            # Squeeze-and-Excite
            # SELayer(hidden_dim) if use_se else nn.Identity(),
            spacial_attention() if use_se else nn.Identity(),
            # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup),
        )

    def forward(self, x):
        if self.identity:
            return x + self.conv(x)
        else:
            return self.conv(x)

class MobileNetV3_large(nn.Module):
    def __init__(self, num_classes=1000, width_mult=1.):
        super(MobileNetV3_large, self).__init__()
        # setting of inverted residual blocks
        self.cfgs_fused = [
            # k, exp,  c,   SE,    NL,  s
            # 112, 112, 16 -> 112, 112, 16
            [3,  16,  16,  False, 'RE', 1],
        ]
        self.cfgs = [
            # 112, 112, 16 -> 56, 56, 24 #2
            [3,  64,  24,  False, 'RE', 2],
            [3,  72,  24,  False, 'RE', 1],

            # 56, 56, 24 -> 28, 28, 40  #3
            [5,  72,  40,  True,  'RE', 2],
            [5,  120, 40,  True,  'RE', 1],
            [5,  120, 40,  True,  'RE', 1],

            # 28, 28, 40 -> 14, 14, 80 #4
            [3,  240, 80,  False, 'HS', 2],
            [3,  200, 80,  False, 'HS', 1],
            [3,  184, 80,  False, 'HS', 1],
            [3,  184, 80,  False, 'HS', 1],

            # 14, 14, 80 -> 14, 14, 112 #5
            [3,  480, 112, True,  'HS', 1],
            [3,  672, 112, True,  'HS', 1],

            # 14, 14, 112 -> 7, 7, 160 #6
            [5,  672, 160, True,  'HS', 2],
            [5,  960, 160, True,  'HS', 1],
            [5,  960, 160, True,  'HS', 1],
        ]

        input_channel = _make_divisible(16 * width_mult, 8)
        # 224,224,3 -> 112,112,16
        layers = [conv_3x3_bn(3, input_channel, 2)]

        block = FusedInvertedResidual
        i = 1
        for k, exp, c, use_se, use_hs, s in self.cfgs_fused:
            output_channel = _make_divisible(c * width_mult, 8)
            exp_size = _make_divisible(exp * width_mult, 8)
            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs))
            input_channel = output_channel
            i = i + 1

        block = InvertedResidual
        for k, exp, c, use_se, use_hs, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, 8)
            exp_size = _make_divisible(exp * width_mult, 8)
            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs, i))
            input_channel = output_channel
            i = i + 1
        self.features = nn.Sequential(*layers)

        self.conv = conv_1x1_bn(input_channel, exp_size)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        output_channel = _make_divisible(1280 * width_mult, 8) if width_mult > 1.0 else 1280
        self.classifier = nn.Sequential(
            nn.Linear(exp_size, output_channel),
            h_swish(),
            nn.Dropout(0.2),
            nn.Linear(output_channel, num_classes),
        )

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.conv(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

    def freeze_backbone(self):
        backbone = [self.features, self.conv]
        for module in backbone:
            print(module)
            for param in module.parameters():
                param.requires_grad = False

    def Unfreeze_backbone(self):
        backbone = [self.features, self.conv]
        for module in backbone:
            print(module)
            for param in module.parameters():
                param.requires_grad = True

class MobileNetV3_small(nn.Module):
    def __init__(self, num_classes=1000, width_mult=1.):
        super(MobileNetV3_small, self).__init__()
        # setting of inverted residual blocks
        self.cfgs_fused = [

            # k, exp,  c,   SE,    NL,  s
            # 112, 112, 16 -> 56, 56, 16 #1
            [3, 16, 16, True, 'RE', 2],

            # 56, 56, 16 -> 28, 28, 24 #2
            [3,  72,  24,  False, 'RE', 2],
            [3,  88,  24,  False, 'RE', 1],
        ]

        self.cfgs = [

            # 28, 28, 24 -> 14, 14, 40 #3
            [5,  96,  40,  True,  'HS', 2],
            [5,  240, 40,  True,  'HS', 1],
            [5,  240, 40,  True,  'HS', 1],

            # 14, 14, 40 -> 14, 14, 48 #4
            [5,  120, 48,  True,  'HS', 1],
            [5,  144, 48,  True,  'HS', 1],

            # 14, 14, 48 -> 7, 7, 96 #5
            [5,  288, 96, True,  'HS', 2],
            [5,  576, 96, True,  'HS', 1],
            [5,  576, 96, True,  'HS', 1],
        ]

        input_channel = _make_divisible(16 * width_mult, 8)
        # 224,224,3 -> 112,112,16
        layers = [conv_3x3_bn(3, input_channel, 2)]

        i = 1
        block = FusedInvertedResidual
        for k, exp, c, use_se, use_hs, s in self.cfgs_fused:
            output_channel = _make_divisible(c * width_mult, 8)
            exp_size = _make_divisible(exp * width_mult, 8)
            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs))
            input_channel = output_channel
            i = i + 1
        block = InvertedResidual
        for k, exp, c, use_se, use_hs, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, 8)
            exp_size = _make_divisible(exp * width_mult, 8)
            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs, i))
            input_channel = output_channel
            i = i + 1
        self.features = nn.Sequential(*layers)

        self.conv = conv_1x1_bn(input_channel, exp_size)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        output_channel = _make_divisible(1024 * width_mult, 8) if width_mult > 1.0 else 1280
        self.classifier = nn.Sequential(
            nn.Linear(exp_size, output_channel),
            h_swish(),
            nn.Dropout(0.2),
            nn.Linear(output_channel, num_classes),
        )

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.conv(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

    def freeze_backbone(self):
        backbone = [self.features, self.conv]
        for module in backbone:
            print(module)
            for param in module.parameters():
                param.requires_grad = False

    def Unfreeze_backbone(self):
        backbone = [self.features, self.conv]
        for module in backbone:
            print(module)
            for param in module.parameters():
                param.requires_grad = True

def mobilenet_v3(pretrained=False, mode='small', **kwargs):
    if mode == 'small':
        model = MobileNetV3_small(**kwargs)
    elif mode == 'large':
        model = MobileNetV3_large(**kwargs)
    else:
        raise NotImplementedError
    if pretrained:
        print('loading siamese model')
        state_dict = torch.load('./siamese/best_epoch_weights.pth')
        model.load_state_dict(state_dict, strict=True)
    return model

if __name__ == "__main__":
    from torchsummary import summary

    # 需要使用device来指定网络在GPU还是CPU运行
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = mobilenet_v3(num_classes=4, mode="small").to(device)
    # print(model)
    summary(model, input_size=(3, 224, 224))

# num_classes=3, mode="small" 1,210,016


